{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Approximate q-learning\n",
    "\n",
    "In this notebook you will teach a lasagne neural network to do Q-learning."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "__Frameworks__ - we'll accept this homework in any deep learning framework. For example, it translates to TensorFlow almost line-to-line. However, we recommend you to stick to theano/lasagne unless you're certain about your skills in the framework of your choice."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%env THEANO_FLAGS = 'floatX=float32'",
    "\n",
    "import os",
    "\n",
    "if type(os.environ.get(\"DISPLAY\")) is not str or len(os.environ.get(\"DISPLAY\")) == 0:",
    "\n",
    "    !bash ../xvfb start",
    "\n",
    "    os.environ['DISPLAY'] = ':1'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import gym",
    "\n",
    "import numpy as np",
    "\n",
    "import pandas as pd",
    "\n",
    "import matplotlib.pyplot as plt",
    "\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env = gym.make(\"CartPole-v0\").env",
    "\n",
    "env.reset()",
    "\n",
    "n_actions = env.action_space.n",
    "\n",
    "state_dim = env.observation_space.shape",
    "\n",
    "\n",
    "plt.imshow(env.render(\"rgb_array\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Approximate (deep) Q-learning: building the network\n",
    "\n",
    "In this section we will build and train naive Q-learning with theano/lasagne"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First step is initializing input variables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import theano",
    "\n",
    "import theano.tensor as T",
    "\n",
    "\n",
    "# create input variables. We'll support multiple states at once",
    "\n",
    "\n",
    "\n",
    "current_states = T.matrix(\"states[batch,units]\")",
    "\n",
    "actions = T.ivector(\"action_ids[batch]\")",
    "\n",
    "rewards = T.vector(\"rewards[batch]\")",
    "\n",
    "next_states = T.matrix(\"next states[batch,units]\")",
    "\n",
    "is_end = T.ivector(\"vector[batch] where 1 means that session just ended\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import lasagne",
    "\n",
    "from lasagne.layers import *",
    "\n",
    "\n",
    "# input layer",
    "\n",
    "l_states = InputLayer((None,)+state_dim)",
    "\n",
    "\n",
    "\n",
    "<Your architecture. Please start with a single-layer network >",
    "\n",
    "\n",
    "\n",
    "# output layer",
    "\n",
    "l_qvalues = DenseLayer( < previous_layer > , num_units=n_actions, nonlinearity=None)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Predicting Q-values for `current_states`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get q-values for ALL actions in current_states",
    "\n",
    "predicted_qvalues = get_output(l_qvalues, {l_states: current_states})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# compiling agent's \"GetQValues\" function",
    "\n",
    "get_qvalues = <compile a function that takes current_states and returns predicted_qvalues >"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# select q-values for chosen actions",
    "\n",
    "predicted_qvalues_for_actions = predicted_qvalues[T.arange(",
    "\n",
    "    actions.shape[0]), actions]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Loss function and `update`\n",
    "Here we write a function similar to `agent.update`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# predict q-values for next states",
    "\n",
    "predicted_next_qvalues = get_output(l_qvalues, {l_states: < theano input with for states > })",
    "\n",
    "\n",
    "\n",
    "# Computing target q-values under",
    "\n",
    "gamma = 0.99",
    "\n",
    "target_qvalues_for_actions = <target Q-values using rewards and predicted_next_qvalues >",
    "\n",
    "\n",
    "# zero-out q-values at the end",
    "\n",
    "target_qvalues_for_actions = (1-is_end)*target_qvalues_for_actions",
    "\n",
    "\n",
    "# don't compute gradient over target q-values (consider constant)",
    "\n",
    "target_qvalues_for_actions = theano.gradient.disconnected_grad(",
    "\n",
    "    target_qvalues_for_actions)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# mean squared error loss function",
    "\n",
    "loss = <mean squared between target_qvalues_for_actions and predicted_qvalues_for_actions >"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# all network weights",
    "\n",
    "all_weights = get_all_params(l_qvalues, trainable=True)",
    "\n",
    "\n",
    "# network updates. Note the small learning rate (for stability)",
    "\n",
    "updates = lasagne.updates.sgd(loss, all_weights, learning_rate=1e-4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Training function that resembles agent.update(state,action,reward,next_state)",
    "\n",
    "# with 1 more argument meaning is_end",
    "\n",
    "train_step = theano.function([current_states, actions, rewards, next_states, is_end],",
    "\n",
    "                             updates=updates)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Playing the game"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "epsilon = 0.25  # initial epsilon",
    "\n",
    "\n",
    "\n",
    "def generate_session(t_max=1000):",
    "\n",
    "    \"\"\"play env with approximate q-learning agent and train it at the same time\"\"\"",
    "\n",
    "\n",
    "    total_reward = 0",
    "\n",
    "    s = env.reset()",
    "\n",
    "\n",
    "    for t in range(t_max):",
    "\n",
    "\n",
    "        # get action q-values from the network",
    "\n",
    "        q_values = get_qvalues([s])[0]",
    "\n",
    "\n",
    "        a = <sample action with epsilon-greedy strategy >",
    "\n",
    "\n",
    "        new_s, r, done, info = env.step(a)",
    "\n",
    "\n",
    "        # train agent one step. Note that we use one-element arrays instead of scalars",
    "\n",
    "        # because that's what function accepts.",
    "\n",
    "        train_step([s], [a], [r], [new_s], [done])",
    "\n",
    "\n",
    "        total_reward += r",
    "\n",
    "\n",
    "        s = new_s",
    "\n",
    "        if done:",
    "\n",
    "            break",
    "\n",
    "\n",
    "    return total_reward"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(100):",
    "\n",
    "\n",
    "    rewards = [generate_session() for _ in range(100)]  # generate new sessions",
    "\n",
    "\n",
    "    epsilon *= 0.95",
    "\n",
    "\n",
    "    print(\"mean reward:%.3f\\tepsilon:%.5f\" % (np.mean(rewards), epsilon))",
    "\n",
    "\n",
    "    if np.mean(rewards) > 300:",
    "\n",
    "        print(\"You Win!\")",
    "\n",
    "        break",
    "\n",
    "\n",
    "    assert epsilon != 0, \"Please explore environment\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Video"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "epsilon = 0  # Don't forget to reset epsilon back to initial value if you want to go on training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# record sessions",
    "\n",
    "import gym.wrappers",
    "\n",
    "\n",
    "env = gym.wrappers.Monitor(gym.make(\"CartPole-v0\"),",
    "\n",
    "                           directory=\"videos\", force=True)",
    "\n",
    "sessions = [generate_session() for _ in range(100)]",
    "\n",
    "env.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# show video",
    "\n",
    "from IPython.display import HTML",
    "\n",
    "import os",
    "\n",
    "\n",
    "video_names = list(",
    "\n",
    "    filter(lambda s: s.endswith(\".mp4\"), os.listdir(\"./videos/\")))",
    "\n",
    "\n",
    "HTML(\"\"\"",
    "\n",
    "<video width=\"640\" height=\"480\" controls>",
    "\n",
    "  <source src=\"{}\" type=\"video/mp4\">",
    "\n",
    "</video>",
    "\n",
    "\"\"\".format(\"./videos/\"+video_names[-1]))  # this may or may not be _last_ video. Try other indices"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
